#!/usr/bin/env python3
"""
Hyperliquid Funding Rate Arbitrage Strategy based on Hummingbot Framework

This strategy implements funding rate arbitrage between spot and perpetual markets
on Hyperliquid using Hummingbot's strategy framework.
"""

import asyncio
import logging
from decimal import Decimal
from typing import Dict, List, Optional, Tuple

from framework.strategy_base import StrategyBase, OrderType, TradeType, PositionSide, OrderFilledEvent, OrderFailedEvent
import asyncio

from config.settings import Config


class HyperliquidFundingArbitrageStrategy(StrategyBase):
    """
    Funding Rate Arbitrage Strategy for Hyperliquid
    
    This strategy captures funding rate differences by:
    1. Going long spot and short perpetual when funding rate is positive
    2. Going short spot and long perpetual when funding rate is negative
    """
    
    def __init__(self, connectors: Dict):
        """
        Initialize the funding arbitrage strategy
        
        Args:
            connectors: Dictionary of exchange connectors
        """
        super().__init__(connectors)
        
        # Load configuration
        self._config = Config()
        
        # Strategy parameters
        self.trading_pair = self._config.TRADING_PAIR or "ETH-USD"
        self.position_size = Decimal(str(self._config.POSITION_SIZE or 1.0))
        self.min_funding_rate = Decimal(str(self._config.MIN_FUNDING_RATE or 0.0001))  # 0.01%
        self.max_position_size = Decimal(str(self._config.MAX_POSITION_SIZE or 10.0))
        self.rebalance_threshold = Decimal(str(self._config.REBALANCE_THRESHOLD or 0.05))  # 5%
        
        # Position tracking
        self._spot_position = Decimal("0")
        self._perp_position = Decimal("0")
        self._target_position = Decimal("0")
        
        # Market data
        self._current_funding_rate: Optional[Decimal] = None
        self._last_funding_check = 0
        self.funding_check_interval = 60  # Check every minute
        
        # Logger
        self._logger = logging.getLogger(__name__)
        
        self._logger.info(f"Initialized Hyperliquid Funding Arbitrage Strategy for {self.trading_pair}")
        self._logger.info(f"Min funding rate: {self.min_funding_rate}, Position size: {self.position_size}")
    
    @property
    def spot_connector(self):
        """Get the spot connector"""
        return self.connectors["hyperliquid_spot"]
    
    @property
    def perp_connector(self):
        """Get the perpetual connector"""
        return self.connectors["hyperliquid_perpetual"]
    
    def on_tick(self):
        """
        Main strategy logic called on each tick
        """
        try:
            # Check if connectors are ready
            if not (self.spot_connector.ready and self.perp_connector.ready):
                self._logger.debug("Connectors not ready, skipping tick")
                return
            
            # Update funding rate
            current_time = self.current_timestamp
            if (current_time - self._last_funding_check) > self.funding_check_interval:
                asyncio.create_task(self._update_funding_rate())
                self._last_funding_check = current_time
            
            # Update positions
            self._update_positions()
            
            # Execute arbitrage logic
            self._execute_arbitrage()
            
        except Exception as e:
            self._logger.error(f"Error in funding arbitrage strategy tick: {e}")
    
    async def _update_funding_rate(self):
        """Update the current funding rate from Hyperliquid API"""
        try:
            # Get funding rate from perpetual connector
            # This would typically involve calling the Hyperliquid API
            # For now, we'll use a placeholder implementation
            
            # In a real implementation, you would call:
            # funding_info = await self.perp_connector.get_funding_info(self.trading_pair)
            # self._current_funding_rate = funding_info.get("funding_rate")
            
            # Placeholder: simulate funding rate
            import random
            self._current_funding_rate = Decimal(str(random.uniform(-0.001, 0.001)))
            
            self._logger.info(f"Updated funding rate: {self._current_funding_rate}")
            
        except Exception as e:
            self._logger.error(f"Error updating funding rate: {e}")
    
    def _update_positions(self):
        """Update current positions from connectors"""
        try:
            # Get spot position (balance)
            spot_balance = self.spot_connector.get_balance(self.trading_pair.split("-")[0])
            self._spot_position = spot_balance if spot_balance else Decimal("0")
            
            # Get perpetual position
            perp_position = self.perp_connector.get_position(self.trading_pair)
            if perp_position:
                self._perp_position = perp_position.amount if perp_position.position_side == PositionSide.LONG else -perp_position.amount
            else:
                self._perp_position = Decimal("0")
            
        except Exception as e:
            self._logger.error(f"Error updating positions: {e}")
    
    def _execute_arbitrage(self):
        """Execute funding rate arbitrage logic"""
        if not self._current_funding_rate:
            return
        
        # Calculate target position based on funding rate
        target_position = self._calculate_target_position()
        
        # Check if rebalancing is needed
        if self._should_rebalance(target_position):
            self._rebalance_positions(target_position)
    
    def _calculate_target_position(self) -> Decimal:
        """Calculate target position based on funding rate"""
        if abs(self._current_funding_rate) < self.min_funding_rate:
            # Funding rate too small, close positions
            return Decimal("0")
        
        if self._current_funding_rate > 0:
            # Positive funding rate: long spot, short perp
            return self.position_size
        else:
            # Negative funding rate: short spot, long perp
            return -self.position_size
    
    def _should_rebalance(self, target_position: Decimal) -> bool:
        """Check if positions need rebalancing"""
        current_net_position = self._spot_position - self._perp_position
        position_diff = abs(current_net_position - target_position)
        
        # Rebalance if difference is significant
        return position_diff > self.rebalance_threshold
    
    def _rebalance_positions(self, target_position: Decimal):
        """Rebalance positions to target"""
        try:
            current_net_position = self._spot_position - self._perp_position
            position_adjustment = target_position - current_net_position
            
            if abs(position_adjustment) < Decimal("0.001"):
                return  # No significant adjustment needed
            
            self._logger.info(
                f"Rebalancing positions: current={current_net_position}, "
                f"target={target_position}, adjustment={position_adjustment}"
            )
            
            # Calculate spot and perp adjustments
            spot_adjustment, perp_adjustment = self._calculate_adjustments(position_adjustment)
            
            # Execute spot adjustment
            if abs(spot_adjustment) > Decimal("0.001"):
                self._execute_spot_trade(spot_adjustment)
            
            # Execute perp adjustment
            if abs(perp_adjustment) > Decimal("0.001"):
                self._execute_perp_trade(perp_adjustment)
            
        except Exception as e:
            self._logger.error(f"Error rebalancing positions: {e}")
    
    def _calculate_adjustments(self, total_adjustment: Decimal) -> Tuple[Decimal, Decimal]:
        """Calculate how to split adjustment between spot and perp"""
        # Simple approach: split adjustment equally
        spot_adjustment = total_adjustment / 2
        perp_adjustment = -total_adjustment / 2
        
        return spot_adjustment, perp_adjustment
    
    def _execute_spot_trade(self, amount: Decimal):
        """Execute spot trade"""
        try:
            if amount > 0:
                # Buy spot
                self.buy_with_specific_market(
                    connector_name="hyperliquid_spot",
                    trading_pair=self.trading_pair,
                    amount=abs(amount),
                    order_type=OrderType.MARKET
                )
                self._logger.info(f"Executed spot BUY: {abs(amount)} {self.trading_pair}")
            else:
                # Sell spot
                self.sell_with_specific_market(
                    connector_name="hyperliquid_spot",
                    trading_pair=self.trading_pair,
                    amount=abs(amount),
                    order_type=OrderType.MARKET
                )
                self._logger.info(f"Executed spot SELL: {abs(amount)} {self.trading_pair}")
                
        except Exception as e:
            self._logger.error(f"Error executing spot trade: {e}")
    
    def _execute_perp_trade(self, amount: Decimal):
        """Execute perpetual trade"""
        try:
            if amount > 0:
                # Long perp
                self.buy_with_specific_market(
                    connector_name="hyperliquid_perpetual",
                    trading_pair=self.trading_pair,
                    amount=abs(amount),
                    order_type=OrderType.MARKET
                )
                self._logger.info(f"Executed perp LONG: {abs(amount)} {self.trading_pair}")
            else:
                # Short perp
                self.sell_with_specific_market(
                    connector_name="hyperliquid_perpetual",
                    trading_pair=self.trading_pair,
                    amount=abs(amount),
                    order_type=OrderType.MARKET
                )
                self._logger.info(f"Executed perp SHORT: {abs(amount)} {self.trading_pair}")
                
        except Exception as e:
            self._logger.error(f"Error executing perp trade: {e}")
    
    def did_fill_order(self, event: OrderFilledEvent):
        """Handle order fill events"""
        try:
            self._logger.info(
                f"Order filled: {event.trade_type.name} "
                f"{event.amount} {event.trading_pair} at {event.price}"
            )
            
            # Update positions after fill
            self._update_positions()
            
        except Exception as e:
            self._logger.error(f"Error handling order fill: {e}")
    
    def _calculate_pnl(self) -> Dict[str, Decimal]:
        """Calculate current PnL from arbitrage positions"""
        try:
            # Get current prices
            spot_price = self._get_mid_price("hyperliquid_spot")
            perp_price = self._get_mid_price("hyperliquid_perpetual")
            
            if not (spot_price and perp_price):
                return {"total_pnl": Decimal("0"), "spot_pnl": Decimal("0"), "perp_pnl": Decimal("0")}
            
            # Calculate unrealized PnL (simplified)
            spot_pnl = self._spot_position * spot_price  # Assuming entry at 0
            perp_pnl = -self._perp_position * perp_price  # Opposite position
            total_pnl = spot_pnl + perp_pnl
            
            return {
                "total_pnl": total_pnl,
                "spot_pnl": spot_pnl,
                "perp_pnl": perp_pnl
            }
            
        except Exception as e:
            self._logger.error(f"Error calculating PnL: {e}")
            return {"total_pnl": Decimal("0"), "spot_pnl": Decimal("0"), "perp_pnl": Decimal("0")}
    
    def _get_mid_price(self, connector_name: str) -> Optional[Decimal]:
        """Get mid price from order book"""
        try:
            connector = self.connectors[connector_name]
            order_book = connector.get_order_book(self.trading_pair)
            if order_book:
                return (order_book.get_price(True) + order_book.get_price(False)) / 2
            return None
        except Exception:
            return None
    
    def format_status(self) -> str:
        """Format strategy status for display"""
        lines = []
        lines.append(f"Strategy: Hyperliquid Funding Arbitrage")
        lines.append(f"Trading Pair: {self.trading_pair}")
        lines.append(f"Current Funding Rate: {self._current_funding_rate}")
        lines.append(f"Min Funding Rate Threshold: {self.min_funding_rate}")
        lines.append("")
        lines.append("Positions:")
        lines.append(f"  Spot Position: {self._spot_position}")
        lines.append(f"  Perp Position: {self._perp_position}")
        lines.append(f"  Net Position: {self._spot_position - self._perp_position}")
        lines.append("")
        
        # PnL information
        pnl_info = self._calculate_pnl()
        lines.append("PnL:")
        lines.append(f"  Total PnL: {pnl_info['total_pnl']}")
        lines.append(f"  Spot PnL: {pnl_info['spot_pnl']}")
        lines.append(f"  Perp PnL: {pnl_info['perp_pnl']}")
        
        return "\n".join(lines)